import numpy as np
from collections import defaultdict
import logging


class DataLoad:
    def __init__(self, file_path = None):
        self.file_path = file_path

    def load_data(self):
        if self.file_path == None:
            logging.error("file_path is None")
            return None
        data = np.loadtxt(self.file_path, delimiter="\t")
        return data


class PointData:
    def __init__(self, file_path = None):
        self.loader = DataLoad(file_path)
        # idx, x, y, type
        self.data = self.loader.load_data()

    def anchor_data(self):
        #return idx, x, y
        return self.data[self.data[:, 3] == 1][:,:3]
    
    def anchor_data_dict(self):
        #return {idx: [x, y]}
        tmp_data = self.data[self.data[:, 3] == 1]
        return dict(
            zip(tmp_data[:, 0], tmp_data[:, 1:3])
        )

    def unknown_data(self):
        #return idx, x, y
        return self.data[self.data[:, 3] == 0][:,:3]
    
    def unknown_data_dict(self):
        #return {idx: [x, y]}
        tmp_data = self.data[self.data[:, 3] == 0]
        return dict(
            zip(tmp_data[:, 0], tmp_data[:, 1:3])
        )
    
    def get_all_points_name(self):
        # return array([idx...])
        return np.hstack(self.data[:, 0])

    def get_point(self, idx):
        # return array([x, y])
        return self.data[self.data[:, 0] == idx][0][1:3]
    def get_points(self, idxs):
        # return array([array([x, y]), ...])
        return np.array([self.get_point(idx) for idx in idxs])

class GraphData:
    def __init__(self, file_path = None):
        self.loader = DataLoad(file_path)

        def generate_graph():
            # u, v, d
            data = self.loader.load_data()
            if data is None:
                return None
            graph = defaultdict(dict)
            for u, v, d in data:
                graph[u][v] = d
                graph[v][u] = d
            return graph

        self.graph = generate_graph()

    def get_distance(self, u, v):
        # return u 到 v 的距离
        return self.graph[u][v]
    def get_neighbors(self, u):
        return self.graph[u].keys()

